from langgraph.graph import StateGraph, END
from typing import TypedDict

# 1. 定义状态
class AgentState(TypedDict):
  question: str
  context: list[str]
  answer: str
  iterations: int

# 2. 定义节点
def retrieve_node(state: AgentState):
    # 模拟检索文档
    state["context"] = ["文档1内容", "文档2内容"]
    return state

def generate_node(state: AgentState):
    # 模拟生成答案
    state["answer"] = f"基于上下文生成的答案，检索了{len(state['context'])}个文档"
    state["iterations"] += 1
    return state

def should_continue_node(state: AgentState):
    # 条件逻辑，迭代次数超过2则停止
    if state["iterations"] >= 2:
      return "end"
    else:
        return "continue"

# 3. 构建图
builder = StateGraph(AgentState)
builder.add_node("retrieve", retrieve_node)
builder.add_node("generate", generate_node)

# 4. 设置入口、边
builder.set_entry_point("retrieve")
# 添加边
builder.add_edge("retrieve", "generate")
# 添加条件边，generate 出来后，进入should_continue_node判断
builder.add_conditional_edges(
  "generate",
  should_continue_node,
  {
    "continue": "retrieve", # 判断返回continue，进入 retrieve 节点
    "end": END
  }
)

# 5. 编译、运行
graph = builder.compile()
initial_state = {
  "question": "我的问题",
  "context": [],
  "answer": "",
  "iterations": 0
}
result = graph.invoke(initial_state)
print(result)